import cv2
import numpy as np

if __name__ == '__main__':
    # 1. 图片输入
    image_np = cv2.imread('document.jpg')

    # 2. 坐标选取（手动标记文档四个角点，按顺时针顺序：左上、右上、右下、左下）
    pts1 = np.float32([
        [50, 20],    # 左上
        [420, 120],  # 右上
        [500, 600],  # 右下
        [40, 670]    # 左下
    ])

    # 将float类型的坐标转换为int类型，用于绘制直线
    pts1_int = pts1.astype(np.int32)

    # 在原图上绘制ROI区域边框，方便验证坐标
    image_line = image_np.copy()
    # 绘制四边形边框（左上→右上→右下→左下→左上）
    cv2.line(image_line, tuple(pts1_int[0]), tuple(pts1_int[1]), (0, 0, 255), 2, cv2.LINE_AA)
    cv2.line(image_line, tuple(pts1_int[1]), tuple(pts1_int[2]), (0, 0, 255), 2, cv2.LINE_AA)
    cv2.line(image_line, tuple(pts1_int[2]), tuple(pts1_int[3]), (0, 0, 255), 2, cv2.LINE_AA)
    cv2.line(image_line, tuple(pts1_int[3]), tuple(pts1_int[0]), (0, 0, 255), 2, cv2.LINE_AA)

    # 3. 定义透视变换的目标矩形（文档矫正后的尺寸，可根据需要调整）
    # 这里使用文档的大致宽高比例（示例：宽600，高800，可根据实际内容调整）
    target_width = 600
    target_height = 800
    pts2 = np.float32([
        [0, 0],                      # 左上
        [target_width, 0],           # 右上
        [target_width, target_height],  # 右下
        [0, target_height]           # 左下
    ])

    # 4. 计算透视变换矩阵并执行变换
    M = cv2.getPerspectiveTransform(pts1, pts2)
    warped = cv2.warpPerspective(
        image_np,
        M,
        (target_width, target_height),
        flags=cv2.INTER_LINEAR,
        borderMode=cv2.BORDER_REPLICATE
    )

    # 5. 图像增强：强化黑白对比+保留红色
    # 5.1 转换为灰度图并二值化（让白色更白，黑色更黑）
    gray = cv2.cvtColor(warped, cv2.COLOR_BGR2GRAY)
    thresh = cv2.adaptiveThreshold(
        gray, 255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,  # 自适应高斯阈值，抗光照不均
        cv2.THRESH_BINARY,
        blockSize=11,  # 阈值计算的区域大小
        C=2            # 常数调整（减小C会让黑色区域更宽）
    )

    # 5.2 提取红色区域（HSV颜色空间分离）
    hsv = cv2.cvtColor(warped, cv2.COLOR_BGR2HSV)
    # 红色HSV范围（覆盖低饱和和高饱和红色）
    lower_red1 = np.array([0, 50, 50])
    upper_red1 = np.array([10, 255, 255])
    lower_red2 = np.array([170, 50, 50])
    upper_red2 = np.array([180, 255, 255])
    mask_red1 = cv2.inRange(hsv, lower_red1, upper_red1)
    mask_red2 = cv2.inRange(hsv, lower_red2, upper_red2)
    red_mask = cv2.bitwise_or(mask_red1, mask_red2)  # 合并红色掩码

    # 5.3 保留红色区域（从原图提取红色部分）
    red_region = cv2.bitwise_and(warped, warped, mask=red_mask)

    # 5.4 合并二值化图像和红色区域（黑白底+红色内容）
    result = cv2.cvtColor(thresh, cv2.COLOR_GRAY2BGR)  # 二值化图转BGR通道
    result = cv2.bitwise_or(result, red_region)  # 叠加红色区域

    # 6. 噪声过滤（中值滤波去除细小噪点）
    result = cv2.medianBlur(result, ksize=3)

    # 7. 显示所有结果
    cv2.imshow('Original with ROI', image_line)  # 带ROI标记的原图
    cv2.imshow('Warped (Perspective Corrected)', warped)  # 透视矫正后图像
    cv2.imshow('Scanned Result (Enhanced)', result)  # 增强后的扫描效果

    # 保存结果
    cv2.imwrite('scanned_result.jpg', result)

    cv2.waitKey(0)
    cv2.destroyAllWindows()